
with ada.text_io; use ada.text_io;
with ada.float_text_io;
use ada.float_text_io;
with ada.numerics.generic_elementary_functions;
with ada.containers.indefinite_vectors;

procedure leet4_median is
    type sequence is array (positive range<>) of integer;

    package foo is new
    ada.numerics.generic_elementary_functions(float);

    package asdf is new ada.containers.indefinite_vectors
        (index_type => positive, element_type => sequence);
    use asdf;

    procedure print(v: in sequence) is
    begin
        for element of v loop
            put(element'image);
        end loop;
        new_line;
    end;

    function partial_merge 
        (a: in sequence; b: in sequence; limit: in integer) 
    return sequence is
        n : integer := a'length + b'length;
        merged :sequence(1 .. limit);
        i  : integer := 1;
        ca : integer := 1;
        cb : integer := 1;
        ae : integer;
        be : integer;
    begin
        for i in 1 .. limit loop
            if ca <= a'length then
                ae := a(ca);
            else
                ae := integer'last;
            end if;
            if cb <= b'length then
                be := b(cb);
            else
                be := integer'last;
            end if;
            if ae < be then
                merged(i) := ae;
                ca := ca + 1; 
            else
                merged(i) := be;
                cb := cb + 1;
            end if;
        end loop;
        return merged;
    end;

    function get_median_even
        (a: in sequence; b: in sequence) return float is
        n : integer := a'length + b'length;
        m1 : integer := n/2;
        m2 : integer := (n/2)+1;
        c : sequence := partial_merge(a, b, m2);
    begin
        put("merged: ");
        print(c);
        return float(c(m1)+c(m2))/2.0;
    end;

    function get_median_odd
        (a: in sequence; b: in sequence) return float is
        n : integer := a'length + b'length;
        m : integer := (n/2)+1; -- index of median
        c : sequence := partial_merge(a, b, m);
    begin
        put("merged: ");
        print(c);
        return float(c(m));
    end;

    -- first solution
    function get_median
        (a: in sequence; b: in sequence) return float is
        n : integer := a'length + b'length;
    begin
        if n mod 2 = 1 then
    -- if my length is 7, then i need the 4th index
            return get_median_odd(a,b);
        else
        -- if my length is 4 then i need 2+2+1/2
        -- if my length is 6 then i need 3+3+1/2
            return get_median_even(a,b);
        end if;
    end;

    function get_median_optimal
        (a : in sequence; b : in sequence) return float is
        n : integer := a'length;
        m : integer := b'length;
        low  : integer := 1;
        high : integer := a'length;
        mid_a : integer;
        mid_b : integer;
        left_a  : integer;
        right_a : integer;
        left_b  : integer;
        right_b : integer;
        answer : float;
        iterations : integer := 0;
        target : float := foo.log(float(m+n));
    begin
        if n < m then
            return get_median_optimal(b, a);
        end if;
        while low <= high loop
            mid_a := (low+high)/2; -- 6/2 = 3
            mid_b := ((1+m+n)/2) - mid_a; -- 10/2 - 3 = 2
            --put_line("mid (a,b): " & mid_a'image & mid_b'image);
            if mid_a <= 0 then
                left_a := integer'first;
            else -- mid_a < a'length then
                left_a := a(mid_a);
            end if;
            if mid_a < a'length then
                right_a := a(mid_a+1);
            else
                right_a := integer'last;
            end if;
            -- b is the shorter array - avoid out of bounds
            if mid_b <= 0 then
                left_b := integer'first;
            elsif mid_b < b'length then
                left_b := b(mid_b);
            elsif mid_b >= b'length then
                left_b := b(b'last);
            end if;
            if mid_b < 0 then
                right_b := b(b'first);
            elsif mid_b < b'length then
                right_b := b(mid_b+1);
            else
                --right_b := b(b'last);
                right_b := integer'last;
            end if;
            --put_line("left (a,b): " & left_a'image & left_b'image);
            --put_line("right (a,b): " & right_a'image & right_b'image);
            -- both left values are less than both right values
            if (left_a <= right_b) 
            and (left_b <= right_a) then
                if ((n + m) mod 2) = 0 then
                    answer := float(
                        integer'max(left_a, left_b) 
                        + integer'min(right_a, right_b)
                    ) / 2.0;
                    exit;
                else
                    answer := float(integer'max(left_a, left_b));
                    exit;
                end if;
            elsif (left_a > right_b) then
                high := mid_a - 1;
            elsif (left_b > right_a) then
                low := mid_a + 1;
            end if;
            iterations := iterations +1;
        end loop;
        put("iterations: " & iterations'image);
        put("; target: O(ln(m+n)) = " & target'image);
        if float(iterations) <=  target then
            put(" ✓");
        end if;
        new_line;
        return answer;
    end;

    function test 
    (x : in sequence; y : in sequence) return boolean is
        s : float := get_median(x,y);
        t : float := get_median_optimal(x,y);
    begin
        put("median =");
        put(s, 4, 1, 0); -- fore, aft, exp
        new_line;
        if s = t then
            return true;
        else
            put_line("result 1: " & float'image(s));
            put_line("result 2: " & float'image(t));
        end if;
        return false;
    end;

    procedure main is
        tests_a : asdf.vector;
        tests_b : asdf.vector;
        pass : integer := 0;
    begin
        tests_a.append((1,3));
        tests_a.append((1,2));

        tests_a.append((1,2,3,4,10));
        tests_a.append((1, 2, 3, 4, 10, 12,13,14,15,16,20,27,30));
        tests_a.append((1,2,3,4,5,6,7,8,9,10,11,12,13,14,15));
        tests_a.append((1,2,3,4));
        tests_a.append(((1,2,3,4,5,6)));
        tests_a.append((2,3,5,7,9));

        tests_b.append((1 => 2));
        tests_b.append((3,4));
        
        tests_b.append((5, 6, 11));
        tests_b.append((5, 6, 11));
        tests_b.append((1 => 17));
        tests_b.append((1 => 1));
        tests_b.append((1 => 0));
        tests_b.append((1,4,6,8,10));

        tests_a.append((1,1,1,1,22,40));
        tests_b.append((1,36,44,62,77,77,77,77));
        for i in 1..integer(tests_a.length) loop
            put_line("----------------------");
            put("a = ");
            print(tests_a(i));
            put("b = ");
            print(tests_b(i));
            if test(tests_a(i), tests_b(i)) then
                put_line("test case #" & i'image & " passed");
                pass := pass +1;
            end if;
        end loop;
        new_line;
        put_line(pass'image 
            & " out of" 
            & tests_a.length'image 
            & " tests pass" );
    end;

begin
    put_line("This is leetcode challenge #4");
    put_line("https://leetcode.com/problems/median-of-two-sorted-arrays/");
    main;
end;

